import numpy as np
from ase.io.trajectory import Trajectory
from ase.visualize import view
from glob import glob
sys.path.append('/media/rbnfiles/ruco/repos/atomistic-analysis/src')
from opendata import data
from pymatgen.io.ase import AseAtomsAdaptor
import pandas as pd
class data:
    def __init__(self):
        self.path  = "/media/rbnfiles/dft/pts2/"
        self.folders = []
        self.files = []
        self.complete_folders= []
        self.complete_files = []
        for folder in glob(self.path+"*"):
            self.folders.append(folder.split('/')[-1])
            self.complete_folders.append(folder)
        for folder in  self.complete_folders:
            for filesinto in glob(folder+'/*.traj'):
                self.complete_files.append(filesinto)

    def get_list_samples(self,folder):
        npath = self.path+'/'+folder
        print(folder)
        self.folder_samples=[]
        for folderi in self.complete_folders:
            if folder in folderi:
                for filesinto in glob(folderi+'/*.traj'):
                    self.folder_samples.append(filesinto.split(self.path)[-1].split('/')[-1])
        return self.folder_samples


    def return_data(self,file):
        for files in self.complete_files:
            if file in files:
                self.select_file = file
                break
        return self.select_file
data().folders

['2x2supercellPtS2', '3x3supercellPtS2', 'teststruct']

files = !ls ../../../../dft/pts2/2x2supercellPtS2/strain/Aniso/*.traj

df = pd.DataFrame([i.split('/')[-1].split('.traj')[0] for i in files],columns=['File'])

no = 0
display(df)
atoms = Trajectory(files[no])
allatoms=[Trajectory(files[i]) for i in range(0,9) if i!= 3]
alllabels=[df.iloc[i][0] for i in range(0,9)if i!=3]
structure = AseAtomsAdaptor.get_structure(atoms[0])
count=0
for i in atoms:
    count+=1
count

	File
0	Opt_VacPt_PtS2_-0.01
1	Opt_VacPt_PtS2_0.01
2	Opt_VacPt_PtS2_-0.02
3	Opt_VacPt_PtS2_0.02
4	Opt_VacPt_PtS2_-0.03
5	Opt_VacPt_PtS2_-0.04
6	Opt_VacPt_PtS2_-0.05
7	Opt_VacPt_PtS2_-0.06
8	Opt_VacPt_PtS2_-0.07
9	Opt_VacPt_PtS2_0.0
10	VacPt_Opt_PtS2

35

import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib.animation as animation
from ase.visualize.plot import plot_atoms
from IPython.display import display, clear_output,HTML
plt.rcParams["animation.html"] = "jshtml"
plt.rcParams["font.family"] = "serif"

ncols = 4
nrows=(df.size/ncols) if isinstance(df.size/ncols, int)==False else int((df.size/ncols)+1)
nrows=2

cm = 1/2.54
px = 1/plt.rcParams['figure.dpi']  # pixel in inches

fig, ax = plt.subplots(figsize=(8,5),frameon=True,facecolor='black')


def trajs(frames):
    ax.clear()
    ax.set_axis_off()
    plot_atoms(atoms[frames],ax,radii=0.3,)
    return fig

def animate(frame):
    clear_output(wait=True)
    print("Run Frame No. :",frame)
    return trajs(frame)

plt.tight_layout()
plt.show()
anim = animation.FuncAnimation(fig,animate,frames=np.arange(0,10),)
anim

Run Frame No. : 9

Once Loop Reflect

ncols = 3
nrows=(len(allatoms)/ncols) if isinstance(len(allatoms)/ncols, int)==False else int((len(allatoms)/ncols)+1)
nrows=3

plt.rcParams["font.family"] = "serif"

fig, ax = plt.subplots(nrows,ncols,figsize=(10,7),frameon=True,facecolor='black')
plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, 
            hspace = 0, wspace = 0)
def alltrajs(frame):
    count=0
    for i in range(nrows):
        for j in range(ncols):
            ax[i,j].clear()
            ax[i,j].set_axis_off()
            if count < len(allatoms):
                ax[i,j].set_title(alllabels[count])
                plot_atoms(allatoms[count][frame],ax[i,j],radii=0.5)
                count+=1
            # else:
            #     fig.delaxes(ax[i,j])
            #     break
    return fig

def animate(frame):
    clear_output(wait=True)
    print("Run Frame No. :",frame)
    return alltrajs(frame)

plt.tight_layout()
plt.show()

#animate(31)
anim = animation.FuncAnimation(fig,animate,frames=np.arange(0,31))
anim

Run Frame No. : 30

Once Loop Reflect

#anim.save('alltrajs.mp4',writer="ffmpeg",fps=10,progress_callback = lambda i, n:print(f'Saving frame {i} of {n}'))
anim.save('alltrajs.gif',writer='imagemagick',fps=10,progress_callback = lambda i, n: print(f'Saving frame {i} of {n}'))

Run Frame No. : 30
Saving frame 30 of 31

files = !ls ../../../../dft/pts2/2x2supercellPtS2/*.traj

df = pd.DataFrame([i.split('/')[-1].split('.traj')[0] for i in files],columns=['File'])

no = 0
display(df)
atoms = Trajectory(files[no])
allatoms=[Trajectory(files[i]) for i in range(0,9) if i!= 3]
alllabels=[df.iloc[i][0] for i in range(0,9)if i!=3]
structure = AseAtomsAdaptor.get_structure(atoms[0])
count=0
for i in atoms:
    count+=1
count

	File
0	Opt_PtS2_VacPt_PW
1	Opt_PtS2_VacPt
2	Opt_PtSe2_VacPt
3	Opt_VacPt_PtS2_-0.01
4	Opt_VacPt_PtS2_0.01
5	Opt_VacPt_PtS2_-0.02
6	Opt_VacPt_PtS2_-0.03
7	Opt_VacPt_PtS2_-0.04
8	Opt_VacPt_PtS2_-0.05
9	Opt_VacPt_PtS2_-0.06
10	Opt_VacPt_PtS2_0.0
11	PtS2_2x2SupvacancyPt
12	PtS2_VacPt_magn
13	PtSe2_2x2SupvacancyPt
14	uCellPtS2
15	VacPt_Opt_PtS2

36

allatoms[3][0]

Atoms(symbols='PtS2PtS2PtS4', pbc=[True, True, False], cell=[[6.138734472185615, 3.6157999999999997, 0.0], [-6.1387344721856145, 3.6158, 0.0], [0.0, 0.0, 18.127]], initial_magmoms=..., calculator=SinglePointCalculator(...))